import numpy as np
from . import LocalSearchOperator, register_operator


class GeneralSwapOperator(LocalSearchOperator):
    """通用交换(General Swap)局部搜索算子"""

    def __init__(self):
        self.name = "general_swap"
        self.supported_problems = ["tsp"]

    def generate_neighborhood(self, solution, problem_instance):
        """生成通用交换邻域"""
        neighborhood = []
        n = len(solution)

        # 生成所有可能的单交换对
        for i in range(n):
            for j in range(i + 1, n):
                new_solution = solution.copy()
                # 交换位置i和j的城市
                new_solution[i], new_solution[j] = new_solution[j], new_solution[i]
                neighborhood.append(new_solution)

        return neighborhood

    def efficient_general_swap(self, solution, problem_instance, intensity=10):
        """高效的通用交换实现，直接计算成本变化"""
        n = len(solution)
        current_solution = solution.copy()
        current_cost = problem_instance.evaluate(current_solution)
        improved = True
        iterations = 0

        while improved and iterations < intensity:
            improved = False

            # 尝试多个随机交换
            for _ in range(min(100, n * (n - 1) // 2)):  # 限制尝试次数
                # 随机选择两个不同的位置
                i, j = np.random.choice(n, 2, replace=False)

                # 计算交换前的相关边成本
                prev_i = (i - 1) % n
                next_i = (i + 1) % n
                prev_j = (j - 1) % n
                next_j = (j + 1) % n

                # 确保索引不重复
                affected_indices = set([prev_i, i, next_i, prev_j, j, next_j])

                # 计算交换前的成本
                old_cost = 0
                for idx in affected_indices:
                    next_idx = (idx + 1) % n
                    old_cost += problem_instance.distance_matrix[
                        current_solution[idx]][current_solution[next_idx]]

                # 执行交换
                new_solution = current_solution.copy()
                new_solution[i], new_solution[j] = new_solution[j], new_solution[i]

                # 计算交换后的成本
                new_cost = 0
                for idx in affected_indices:
                    next_idx = (idx + 1) % n
                    new_cost += problem_instance.distance_matrix[
                        new_solution[idx]][new_solution[next_idx]]

                # 计算成本变化
                cost_change = new_cost - old_cost

                # 如果成本降低，则接受交换
                if cost_change < 0:
                    current_solution = new_solution
                    current_cost += cost_change
                    improved = True
                    break  # 找到一个改进就跳出内层循环

            iterations += 1

        improvement = problem_instance.evaluate(solution) - current_cost
        return current_solution, improvement

    def __call__(self, solution, problem_instance, intensity=10):
        """应用通用交换局部搜索"""
        if not self.is_applicable(problem_instance):
            return solution, 0

        # 使用高效实现
        return self.efficient_general_swap(solution, problem_instance, intensity)


# 注册算子
register_operator("general_swap", GeneralSwapOperator())